#!/usr/bin/python3
# -*- coding: utf-8 -*-

import os
import sys
import cv2
import tqdm
import json
import argparse
import numpy as np
from terminaltables import AsciiTable

from module.annotations import load_annotations
from module.detections import load_detections
from module.common import calculate_iou, calculate_iou2
from module.common import get_dalei_categories, trans2dalei
from module.common import draw_objects
from module.common import get_name_id_map

class Precisoin(object):
    def __init__(self, args):
        self.args = args
        if args.output_folder:
            if os.path.exists(args.output_folder):
                print("please remove output folder {}".format(args.output_folder))
                sys.exit(0)
            else:
                os.makedirs(args.output_folder)

    def __del__(self):
        pass

    def convert2result(self, annotations, detection_file):
        det_out = []
        object_id = 1
        for ann in annotations:
            for an in annotations[ann]["gt"]:
                item = {
                    "id": object_id,
                    "path": ann,
                    "type": an["category"],
                    "score": 0.99,
                    "xmin": an["bbox"][0],
                    "ymin": an["bbox"][1],
                    "xmax": an["bbox"][2],
                    "ymax": an["bbox"][3]
                }
                det_out.append(item)
                object_id += 1
        with open(detection_file, "w") as file:
            json.dump(det_out, file)
    
    def load_data(self):
        print("using category file {}".format(self.args.categories_file))
        with open(self.args.categories_file) as file:
            categories = json.load(file)
            assert(len(categories) > 0)
        
        for cat in categories:
            categories[cat]["dtn"] = 0
            categories[cat]["gtn"] = 0
            categories[cat]["del"] = 0
            categories[cat]["matched"] = 0
        file_id_map, id_file_map, id_categorie_map = get_name_id_map(self.args.imageset)
        # print (file_id_map)
        # print(id_file_map)
        # print(id_categorie_map)

        categories, annotations = load_annotations(categories, self.args.annotations_folder, file_id_map)
        self.categories, self.detections = load_detections(categories, annotations, self.args.detection_file,
                                                           id_file_map,id_categorie_map)
        

    def get_matched_pair(self, annos, dets):
        matched_index = {}
        output_dets = []
        used_det_index = []
        for det_index, det in enumerate(dets):
            # 统计大于阈值的数量
            if det["score"] < self.categories[det["category"]]["thresh"]:
                continue
            else:
                self.categories[det["category"]]["dtn"] += 1
                output_dets.append(det)

        for ann_index, ann in enumerate(annos):
            iou_thresh = self.categories[ann["category"]]["iou"]
            matched_det_index = -1
            for det_index, det in enumerate(output_dets):
                if det["category"] != ann["category"]: # 跳过不匹配的类别
                    continue
                if det_index in used_det_index:
                    continue
                iou = calculate_iou(ann["bbox"], det["bbox"])
                if iou > iou_thresh:  # 替换为与标注最匹配的框
                    iou_thresh = iou
                    matched_det_index = det_index
            if matched_det_index != -1:  # 存在标注文件
                matched_index[ann_index] = matched_det_index
                used_det_index.append(matched_det_index)
        return matched_index, output_dets

    def get_matched_count(self, annos, dets, matched_index):
        if len(matched_index) == 0:
            return

        for ann_index, ann in enumerate(annos):
            if ann_index not in matched_index:
                continue
            self.categories[ann["category"]]["matched"] += 1

    def get_valid_flag(self, annos, dets):
        annos_cls = {}
        for ann in annos:
            if ann["category"] not in annos_cls:
                annos_cls[ann["category"]] = 1
            else:
                annos_cls[ann["category"]] += 1
        dets_cls = {}
        for det in dets:
            if det["category"] not in dets_cls:
                dets_cls[det["category"]] = 1
            else:
                dets_cls[det["category"]] += 1

        # 只要存在超过2/3倍的情况，认为该图没有检测结果
        flag = True
        for cls in annos_cls:
            if cls in dets_cls:
                #limit_num = self.get_limit_number(cls)
                limit_num = self.get_type_limit(cls)
                if dets_cls[cls] > limit_num * annos_cls[cls]:
                    flag = False
                    self.categories[cls]["del"] += 1

        return flag

    def get_type_limit(self,cls):
        return self.categories[cls]["max_match"]

    def get_precision(self):
        """
        按照类别统计精度
        """
        #print(self.categories)
        noalarm_image_number = 0
        detected_image_number = 0
        miss_alarm, false_alarm = 0, 0

        print("calculating...")
        for image_name in tqdm.tqdm(self.detections):
            annos = self.detections[image_name]["gt"]
            dets = self.detections[image_name]["dt"]

            # 将标注结果与检测结果进行匹配
            matched_index, dets = self.get_matched_pair(annos, dets)

            flag = self.get_valid_flag(annos, dets)#2倍框限制验证，不存在标注时，flag=True
            if not flag:# 超过2倍框认为是误检
                false_alarm += 1
            else:# 计算匹配成功的个数，用于计算准确率
                self.get_matched_count(annos, dets, matched_index)

            if len(dets) > 0: #存在检测结果的图像数量
                detected_image_number += 1
    
            if len(annos) == 0: # 不存在标注的图像数量
                noalarm_image_number += 1
            
            if len(annos) == 0 and len(dets) > 0:  # 不存在标注，但是有检测结果
                false_alarm += 1
 
            if len(annos) > 0 and len(dets) == 0: # 没有检测结果就认为是漏报
                miss_alarm += 1

            if self.args.output_folder:
                image_abs_name = os.path.join(self.args.images_folder, image_name+".jpg")
                image = cv2.imread(image_abs_name)
                try:
                    shape = image.shape
                    image = draw_objects(image, annos, dets, matched_index, self.categories)
                    #image_abs_name = os.path.join(self.args.output_folder, output_subfolder, image_name+".jpg")
                    image_abs_name = os.path.join(self.args.output_folder,  image_name + ".jpg")
                    image_folder = os.path.split(image_abs_name)[0]
                    if not os.path.exists(image_folder):
                        os.makedirs(image_folder)
                    cv2.imwrite(image_abs_name, image)
                except AttributeError:
                    print("load image {} failed".format(image_name))
                    sys.exit()

        table_data = [['classname', 'thresh', 'iou', 'weight', 'max_match', 'gtn', 'dtn', 'matched', 'delete', 'percentage']]
        total_gt_num = 0
        total_matched_num = 0
        for cat_name in self.categories:
            if cat_name == "FalseAlarm" or cat_name == "MissAlarm":
                continue
            total_gt_num += self.categories[cat_name]["gtn"]
            total_matched_num += self.categories[cat_name]["matched"]
            self.categories[cat_name]["percentage"] = self.categories[cat_name]["matched"]/self.categories[cat_name]["gtn"] if self.categories[cat_name]["gtn"] > 0 else 0.0

            row_data = [cat_name, 
                        '{:.2f}'.format(self.categories[cat_name]["thresh"]),
                        '{:.2f}'.format(self.categories[cat_name]["iou"]),
                        '{:.2f}'.format(self.categories[cat_name]["weight"]),
                        self.categories[cat_name]["max_match"],
                        self.categories[cat_name]["gtn"], 
                        self.categories[cat_name]["dtn"],
                        self.categories[cat_name]["matched"],
                        self.categories[cat_name]["del"],
                        '{:.2f}'.format(self.categories[cat_name]["percentage"]*100)]
        table_data.append(row_data)
        table = AsciiTable(table_data)
        table.inner_footing_row_border = True
        table.justify_columns = {0: 'center', 1: 'center', 2: 'center'}
        table.title = "每个类别检出情况"
        print(table.table)

        # 输出得分
        alarm_image_number = len(self.detections) - noalarm_image_number
        score_false = self.categories["FalseAlarm"]["weight"]*false_alarm/detected_image_number
        score_miss = self.categories["MissAlarm"]["weight"]*miss_alarm/alarm_image_number
        acc_score, acc_weight = 0, 1-self.categories["FalseAlarm"]["weight"]-self.categories["MissAlarm"]["weight"]
        for cat_name in self.categories:
            if cat_name == "FalseAlarm" or cat_name == "MissAlarm":
                continue
            acc_score += self.categories[cat_name]["weight"] * self.categories[cat_name]["percentage"]
        score_acc = acc_weight*(1-acc_score)
        score = 1 - score_miss - score_false - score_acc

        table_data = [['task', 'alarm_images/detected_images/gt_objects', 'count', 'percentage', 'score']]
        row_data = ["漏报", 
            alarm_image_number, 
            miss_alarm, 
            '{:.2f}'.format(miss_alarm*100/alarm_image_number) if alarm_image_number>0 else 0,
            '{:.2f}'.format(score_miss*100)]
        table_data.append(row_data)

        row_data = ["误报", 
            detected_image_number, 
            false_alarm, 
            '{:.2f}'.format(false_alarm*100/detected_image_number) if detected_image_number>0 else 0,
            '{:.2f}'.format(score_false*100)]
        table_data.append(row_data)
        
        row_data = ["准确", 
                total_gt_num, 
                total_matched_num, 
                '{:.2f}'.format(total_matched_num*100/total_gt_num),
                '{:.2f}'.format(score_acc*100)]
        table_data.append(row_data)

        row_data = [
            "最终",
            "",
            "",
            "",
            '{:.2f}'.format(score*100)
        ]
        table_data.append(row_data)

        table = AsciiTable(table_data)
        table.inner_footing_row_border = True
        table.justify_columns = {0: 'center', 1: 'center', 2: 'center'}
        table.title = "准确率漏报误报统计表"
        print(table.table)
 

    # def get_dalei_precision(self):
    #     """
    #     按照大类类别统计精度
    #     """
    #     print("#############################DaLei##################################")
    #     print("calculating dalei precision...")
    #
    #     self.categories = get_dalei_categories(self.categories)
    #     #print(self.categories)
    #     self.get_precision(xiaolei_flag=False)